open Graphics
open Hashtbl
open Utils.Pervasive
open Utils.RandSelect
open Syntax
open Binding
open PrintExpr
open Core
open GPUtils



module TwoDRoutines = struct

	type coordinates =  {mutable xPos : int; mutable yPos : int}
	
	
	let coord_of_string  str = 
	let reg =  Str.regexp "(\([0-9]+\),\([0-9]+\))" in
	if Str.string_match reg str 0 then (
		let str1 = Str.matched_group 1 str in
		let str2 = Str.matched_group 2 str in  
		{xPos = int_of_string str1; yPos = int_of_string str2})
		else failwith "Unable to produce coordinates from string" 		
	let string_of_coord co = "(" ^ string_of_int (co.xPos) ^ "," ^ string_of_int (co.yPos) ^ ")"
    let helpMessage str = moveto 0 0; draw_string str;
    (*see 
        - imperative features in http://caml.inria.fr/pub/docs/manual-ocaml/manual003.html
        - mutable fields of record, o'reilly pg 73*)
    
    type twoDSpace = {
		mutable width : int; 
		mutable height : int; 
		mutable numCols : int; 
		mutable numRows : int}
    let makeTwoDSpace w h nC nR = {width=w; height=h; numCols=nC; numRows=nR}
    let stepX space = int_of_float((float space.width) /. (float space.numCols))
    let stepY space = int_of_float ((float space.height) /. (float space.numRows))
    let getNc space = space.numCols
    let getNr space = space.numRows
    let setGrid nC nR space = 
        let assignG() = space.numCols<-nC; space.numRows<-nR
        in
        helpMessage ((string_of_int nC)^"-"^(string_of_int nR));
        if (nC >= 10) && (nC <= 100)&&(nR >= 10)&&(nR <= 100) then assignG() 
                                
                                
    
    let buildCoordinates x y = {xPos = x; yPos=y}
    let _x = function {xPos=x; yPos=_}->x
    let _y = function {xPos=_; yPos=y}->y
    
    let scalProd co1 co2 =  ((co1.xPos - co2.xPos)*(co1.xPos - co2.xPos))+
                            ((co1.yPos - co2.yPos)*(co1.yPos - co2.yPos))
    let coordDistSq co1 co2 = 
        let n = scalProd co1 co2 in
        if n < 0 then (-1*n) else n
    
    let coord_to_xy coord space = (coord.xPos*(stepX space), coord.yPos*(stepY space)) 
    let xy_to_coord x y space = buildCoordinates (x/(stepX space)) (y/(stepY space)) 
    
    
    let is_in_grid x y space = 
        let coord = xy_to_coord x y space in
       (_x coord) >= 0 && (_x coord) < space.numCols 
        && (_y coord) >= 0 && (_y coord < space.numRows) 
       
    
    let assignCoord x y co space= 
        let toroiMove pos surfaceLength = 
        if pos < 0 then surfaceLength-1
        else if pos >= surfaceLength then 0
        else pos
        in
        co.xPos<-(toroiMove x (getNc space));
        co.yPos<-(toroiMove y (getNr space))
    let draw_pixel col pt space=
        set_color col;
        if pt.xPos < space.numCols && pt.yPos < space.numRows then  
            let (x, y)= coord_to_xy pt space in 
            fill_rect x y (stepX space) (stepY space )
        else()
    
    let is_in_square_bound x y = function 
        ({xPos=xP; yPos=yP},w ,h) 
        when (x < xP + w) && (y < yP + h) && (x >= xP)&& (y >= yP)->true
        |_->false
   
    let rec draw_grid space= 
        set_color red;
        let rec draw_grid_vertline = function 
            xCo when xCo <= space.numCols -> 
                moveto (xCo *(stepX space)) 0;
                lineto (xCo *(stepX space)) space.height;
                draw_grid_vertline (xCo+1)
            |_-> ()  
        in  
        let rec draw_grid_horline = function 
            yCo when yCo <= space.numRows -> 
                moveto 0 (yCo * (stepY space));
                lineto ((stepX space) * space.numCols) (yCo * (stepY space));
                draw_grid_horline (yCo+1)
            |_-> ()  
        in  
        draw_grid_vertline 0;
        draw_grid_horline 0
    type twoDVec = {mutable c1 : coordinates; mutable tip : coordinates}
end (* module TwoDRoutines *)

open TwoDRoutines
module AntEnv = struct
	type ant = coordinates * bool ref
	type pheromones = (coordinates, int) t
	type food = (coordinates, int) t
	type antHill =  coordinates * int ref
	type env = ant list  * pheromones * food * antHill
	
	
	let getLeft space co = {
				xPos = (if co.xPos = 0 then space.numCols -1 else co.xPos -1);
				yPos = co.yPos}
	let getRight space co = {
		xPos = (if co.xPos = space.numCols -1 then  0 else co.xPos  + 1);
		yPos = _y co}
	let getUp space co = {
		xPos = _x co;
		yPos =  (if co.yPos = space.numRows -1 then  0 else co.yPos  + 1)}
	let getDown space co = {
		xPos = _x co;
		yPos = (if co.yPos = 0 then  space.numRows  - 1 else co.yPos - 1)}
	
	
	let addPheromone co = 
		function (_, pheros,_,_) -> add pheros {xPos = co.xPos; yPos = co.yPos} 255	
	
	let resetPhero pheros = 
		iter (fun co bright -> 
			if bright < 20 then remove pheros co
			else replace pheros co (bright - 20)) pheros
		
	let is_in_hill co = 
		((co.xPos = 10) || (co.xPos = 11)) &&
		((co.yPos = 10) ||  (co.yPos = 11))
	
	let initEnv space num_food_parts = 
		let randCoord space = 
			{xPos = (Random.int space.numCols); yPos = (Random.int space.numRows)} in
		let hill  = ({xPos = 10; yPos =10}, ref 0) in
		let createFood space size = 
			let food = create size in 
			let co = ref (randCoord space) in
			let i = ref 0 in
			while !i < size do
				if (Hashtbl.mem food !co) or (is_in_hill !co) then ( )
				else (add food !co 1; i := !i + 1);
				match Random.int 5 with
				| 0 -> co := getLeft space !co
				| 1 -> co := getUp space !co
				| 2 -> co := getDown space !co
				| 3 -> co := getRight space !co
				| 4 -> co := randCoord space 		
			done;
			food	in
		let initAntList  = [
		({xPos = 10; yPos =10}, ref false);
		({xPos = 10; yPos =10}, ref false);
		({xPos = 10; yPos =10}, ref false)] in 
		(initAntList, create 10, createFood space num_food_parts, hill)
	

	
	let drawEnv (antList, pheromones, food, hill) space = 
		let drawHill( ) = 
			draw_pixel (rgb 100 100 0) {xPos=10;yPos=10} space;
			draw_pixel (rgb 100 100 0) {xPos=11;yPos=10} space;
			draw_pixel (rgb 100 100 0) {xPos=10;yPos=11} space;
			draw_pixel (rgb 100 100 0) {xPos=11;yPos=11} space in
			
		let drawFood co = 
			draw_pixel (rgb 0 255 0) co space in
		let drawPhero coord bright = 
			draw_pixel (rgb 0 0 bright) coord space in
		let drawAnt co food =
			draw_pixel (rgb (if !food then 255 else 100) 0 0) co space in		

		List.iter (fun (co,food) -> drawAnt co food) antList;
		Hashtbl.iter (fun co _ -> drawFood co) food;
		Hashtbl.iter drawPhero pheromones;
		drawHill( )	
	let lSFtrue ctx = ter (parse_string "Lam X . lam x : X . lam y : X . x" ctx)
	let lSFfalse ctx = ter (parse_string "Lam X . lam x : X . lam y : X . y" ctx)
	
	let rec eval1 (co,carry) ctx t space env_ref res_ref  =
	let same op str = name_of_index ctx op = str in
	match !env_ref with  (antList, pheromones, food, hill) ->
	match t with 
		| TmTApp (t1, t2) ->
			TmTApp (eval1 (co,carry) ctx t1 space env_ref res_ref, t2) 		
		| TmApp(TmApp(TmVar op, arg1), arg2) when same op "contains" ->
			(match arg1,arg2 with
			| TmString strCoord, TmVar obj  ->
				let cellDetected = coord_of_string strCoord in
				(match name_of_index ctx  obj with
				|"ant"-> raise NoRuleApplies
				| "food"-> (if Hashtbl.mem food cellDetected then lSFtrue else lSFfalse) ctx
				| "pheromone" -> (if Hashtbl.mem pheromones cellDetected then lSFtrue else lSFfalse) ctx
				|_->raise NoRuleApplies)
			| TmString strCoord,_->
				TmApp(TmApp(TmVar op, TmString strCoord), eval1 (co,carry) ctx arg2 space env_ref res_ref)	
			| _-> TmApp(TmApp(TmVar op, eval1 (co,carry) ctx arg1 space env_ref res_ref), arg2)	)			
		| TmApp(TmVar op, arg) when same op "detect" ->
			(match arg with 
			| TmVar arg ->
				let _p op = TmString(string_of_coord (op space co)) in
				(match name_of_index ctx arg with
					| "left" -> _p getLeft							
					| "right" -> _p getRight 							
					| "up" -> _p getUp 
					| "down" -> _p getDown 
					| "center" -> _p (fun _ x  -> x)
					| _-> raise NoRuleApplies)
			| _ -> TmApp(TmVar op, eval1 (co,carry) ctx arg space env_ref res_ref)) 	
		| TmApp(TmVar op, arg) when same op "move" ->
			(match arg with 
			|TmVar arg ->
			(match name_of_index ctx arg with 
				| "left" -> co.xPos <- _x (getLeft space co);TmUnit	
				| "right" -> co.xPos <- _x (getRight space co);TmUnit	
				| "up" -> co.yPos <- _y (getUp space co);TmUnit	
				| "down" -> co.yPos <- _y (getDown space co);TmUnit
				| "home" -> 
					if not (is_in_hill co) then	(
						match co.xPos,co.yPos with
						| x, y when x > 11 && x < 20 && y > 11 -> co.xPos <- _x (getLeft space co)
						| x,y when y > 11 && y < 20 ->  co.yPos <- _y (getDown space co)
						| x,y when x < 10 || x > 20 -> co.xPos <- _x (getRight space co)
						| x,y when y < 10 || y > 11 -> co.yPos <- _y (getUp space co)
						| x, y when x > 11 && x < 20 -> co.xPos <- _x (getLeft space co)
						|_-> ( ));
					TmUnit	
				| "center" -> TmUnit	
				| "random" -> 
					TmApp(TmVar op,
					TmVar(index_of_name ctx (List.hd (choose_random ["left";"right";"up";"down"]))))
				| _-> raise NoRuleApplies)
			|_-> TmApp(TmVar op, (eval1 (co,carry) ctx arg space env_ref res_ref)))
		| TmApp (t1, t2) ->
			(try TmApp (eval1 (co,carry) ctx t1 space env_ref res_ref, t2) 
			with NoRuleApplies ->
				TmApp (t1, eval1 (co,carry) ctx t2 space env_ref res_ref)) 		
		| TmVar op when same op "dropPheromone" ->	
			 		addPheromone co !env_ref;
			 		TmUnit
		| TmVar op when same op "pickupFood" ->	
			if  !carry then raise NoRuleApplies
			else (
				if Hashtbl.mem food co then
				(Hashtbl.remove food co;
				carry := true;
				res_ref := !res_ref +. 0.5;
				TmUnit) else raise NoRuleApplies)
		| TmVar op when same op "dropFood" ->
			if not !carry then raise NoRuleApplies
			else (
				if (not (is_in_hill co)) then (
					Hashtbl.add food {xPos = co.xPos; yPos = co.yPos} 1;
					res_ref := !res_ref -. 0.5)
				else res_ref := !res_ref +. 0.5;
				carry := false;
				TmUnit)
		| TmVar op when same op "load" ->
			if !carry then lSFtrue ctx else lSFfalse ctx
		| TmVar op when same op "isHome" ->
			if is_in_hill co then lSFtrue ctx else lSFfalse ctx
		|_->raise NoRuleApplies
	
	
		
	let rec eval ctx t space env_ref ant res_ref =
			(*prf (string_of_term ctx t Alpha ^ "\n");*)
			let t' = normalize ctx (eval1 ant ctx t space env_ref res_ref)
			in match t' with 
			|TmUnit -> ( )
			|_-> eval ctx t' space env_ref ant res_ref
			
			
			
	let rec play ctx ter env_ref space res_ref = 
		match !env_ref with (ants, pheros, food, hill) ->
			List.iter (fun ant -> eval ctx ter space env_ref ant res_ref) ants;
			resetPhero pheros

	let playAnt ?(visual=true) ticker num_food_parts ctx ter  = 
		if visual then (
			open_graph " 840x480";
			set_window_title "Playing ant in sssenv");
		
		let space = makeTwoDSpace 600 600 30 30 in
		let env = ref (initEnv space (int_of_float num_food_parts)) in 
		let evaluation_result = ref 0.0 in
		let goOn = ref true in
		
		while (ticker( ) && !goOn) do
			if visual then 
				(draw_grid space;
				drawEnv !env space);
			(try 
				play ctx ter env space evaluation_result;
				evaluation_result := !evaluation_result +. 0.01; (*Token plus for not crashing*)
				with _ -> evaluation_result := 0.0; goOn := false);
			if visual then (
				let i = ref 10000000 in
				while !i > 0 do i := !i-1 done; 
				clear_graph( ) );
		done;
		prf (" - obtained " ^ string_of_float  !evaluation_result);
		!evaluation_result
end

